import requests
import json
import os
from typing import List, Dict

class GLMService:
    def __init__(self):
        self.api_url = os.getenv('CODEGEEX_API_URL')
        self.api_key = os.getenv('CODEGEEX_API_KEY')
        self.model_id = os.getenv('MODEL_ID')
        self.conversations: Dict[str, List[Dict]] = {}  # 存储多个会话的历史记录

    def chat(self, message: str, conversation_id: str = None) -> str:
        """
        Handle chat messages using GLM model with conversation memory
        Args:
            message (str): User's chat message
            conversation_id (str): Unique identifier for the conversation
        Returns:
            str: Model's response
        """
        if not conversation_id:
            conversation_id = "default"

        # 初始化或获取现有的对话历史
        if conversation_id not in self.conversations:
            self.conversations[conversation_id] = []

        # 添加用户消息到对话历史
        self.conversations[conversation_id].append({
            "role": "user",
            "content": message
        })

        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }

        # 使用完整的对话历史构建请求
        data = {
            'model': self.model_id,
            'messages': self.conversations[conversation_id],
            'temperature': 0.7,
            'max_tokens': 1000
        }

        try:
            response = requests.post(self.api_url, headers=headers, json=data)
            response.raise_for_status()
            result = response.json()
            assistant_message = result['choices'][0]['message']['content']

            # 将助手的回复添加到对话历史
            self.conversations[conversation_id].append({
                "role": "assistant",
                "content": assistant_message
            })

            # 如果对话历史过长，可以进行裁剪
            if len(self.conversations[conversation_id]) > 20:  # 保留最近的20轮对话
                self.conversations[conversation_id] = self.conversations[conversation_id][-20:]

            return assistant_message
        except Exception as e:
            print(f"Error in chat: {str(e)}")
            return "I apologize, but I'm having trouble processing your message right now. Please try again later."

    def clear_conversation(self, conversation_id: str = None) -> bool:
        """
        Clear the conversation history for a specific conversation
        Args:
            conversation_id (str): Unique identifier for the conversation
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            if not conversation_id:
                conversation_id = "default"
            if conversation_id in self.conversations:
                self.conversations[conversation_id] = []
            return True
        except Exception as e:
            print(f"Error clearing conversation: {str(e)}")
            return False


    def generate_test_cases(self, code_content, requirements, module_name):
        """
        Generate test cases using GLM model

        Args:
            code_content (str): Source code content to generate tests for
            requirements (str): Test requirements specified by user
            module_name (str): The name of the module to import in test cases.
            
        Returns:
            list: Generated test cases as a list of strings
        """
        # 从文件路径中提取模块名
        prompt = self._build_prompt(code_content, requirements, module_name)

        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }

        data = {
            'model': self.model_id,
            'messages': [{
                'role': 'user',
                'content': prompt
            }],
            'temperature': 0.7,
            'max_tokens': 2000
        }

        try:
            response = requests.post(self.api_url, headers=headers, json=data)
            response.raise_for_status()

            result = response.json()
            generated_code = result['choices'][0]['message']['content']

            test_cases = self._parse_test_cases(generated_code)
            return test_cases

        except Exception as e:
            print(f"Error calling GLM API: {str(e)}")
            raise

    def _build_prompt(self, code_content, requirements, module_name):
        """
        Builds the prompt to instruct the GLM model to generate test cases.
        
        Args:
            code_content (str): The source code content
            requirements (str): Test requirements
            module_name (str): Name of the module being tested
        """
        prompt_template = """Generate pytest test cases for the following Python code.
Requirements: {requirements}

Source code from {module_name}.py:
{code_content}

Please create pytest test cases following these strict guidelines:

1. Format and Structure:
   - Start with a single 'import pytest' statement
   - Import the necessary classes/functions from {module_name}.
   - Use proper Python class naming conventions (e.g., TestClassName)
   - Include clear docstrings for the test class and each test method
   - Use proper indentation (4 spaces)
   - No empty lines at the start of the file
   - No explanatory comments or documentation at the end

2. Test Implementation:
   - Use setup_method() for test fixture initialization when needed
   - Write independent test methods for each functionality
   - Name test methods with prefix 'test_'
   - Include assertions that clearly validate expected outcomes
   - Use pytest.approx() for floating-point comparisons
   - Use pytest.raises() for testing exceptions
   - Keep test methods focused and concise

3. Best Practices:
   - Test both valid and edge cases
   - Include positive and negative test scenarios
   - Test boundary conditions
   - Use descriptive variable names
   - Make assertions explicit and clear
   - Avoid redundant test cases
   - No print statements or debugging code

4. Important:
   - Use only ASCII characters in the code
   - No Chinese characters or other non-ASCII text
   - No explanatory text or comments in natural language after the test code
   - Include only the actual test code that will be executed

Example format:
import pytest
from module import ClassName

class TestClassName:
    \"\"\"Test cases for ClassName.\"\"\"

    def setup_method(self):
        \"\"\"Setup test fixtures.\"\"\"
        self.instance = ClassName()

    def test_method_name(self):
        \"\"\"Test specific functionality.\"\"\"
        result = self.instance.method(param)
        assert result == expected_value

Generate the test cases now:"""
        
        return prompt_template.format(
            requirements=requirements,
            module_name=module_name,
            code_content=code_content
        )

    def _parse_test_cases(self, generated_code):
        """
        Parse the generated code into a list of test cases, removing explanatory text
        Replace any 'source_code' imports with the actual module name.

        Args:
            generated_code (str): Raw generated code from GLM model

        Returns:
            list: Clean test case code lines
        """
        code = generated_code.replace('```python', '').replace('```', '').strip()
        lines = code.split('\n')
        parsed_lines = ['import pytest']
        in_code_block = False

        for line in lines:
            if not in_code_block and not line.strip():
                continue

            if (line.startswith('import ') or
                line.startswith('from ') or
                line.startswith('def ') or
                line.startswith('@pytest') or
                line.startswith('class ')):
                in_code_block = True

            if in_code_block:
                if (line.startswith('Make sure') or
                    line.startswith('These test') or
                    line.startswith('You can') or
                    line.startswith('This test')):
                    break

                stripped_line = line.strip()
                if (stripped_line and
                    not stripped_line.startswith('# Below is') and
                    not stripped_line.startswith('# Here')):
                    parsed_lines.append(line)

        while parsed_lines and not parsed_lines[-1].strip():
            parsed_lines.pop()

        return parsed_lines